package hku.cs.lunuo.profile;

import hku.cs.lunuo.constant.DirAndFileConstant;
import hku.cs.lunuo.tools.DirTools;
import hku.cs.lunuo.tools.FindClass;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

import org.jdom.Document;
import org.jdom.Element;
import org.jdom.output.XMLOutputter;

import soot.Body;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.Unit;
import soot.Value;
import soot.jimple.IfStmt;
import soot.options.Options;
import soot.tagkit.LineNumberTag;
import soot.toolkits.graph.Block;
import soot.toolkits.graph.ExceptionalBlockGraph;
import soot.util.Chain;

class Cond {
	int blk;
	int lineNum;
	Value cond;
	int startLine;
	int endLine;
	List<Integer> lineLst;

	public Cond(int blk, int lineNum, Value cond) {
		super();
		this.blk = blk;
		this.lineNum = lineNum;
		this.cond = cond;
	}

	public Cond(int blk, int lineNum, Value cond, List<Integer> lineLst) {
		super();
		this.blk = blk;
		this.lineNum = lineNum;
		this.cond = cond;
		this.lineLst = lineLst;
	}

}

public class PredsProfile {
	public static void main(String[] args) {

		String output_dir = null;
		String process_dir = null;
		List<String> set = new ArrayList<String>();
		for (int i = 0; i < args.length; i++) {
			if (args[i].equals("-process-dir"))
				process_dir = args[++i];
			else if (args[i].equals("-output-dir"))
				output_dir = args[++i];
			else
				set.add(args[i]);
		}

		if (output_dir == null)
			return;

		if (process_dir != null) {
			File f = new File(process_dir);
			if (f.exists() && f.isDirectory())
				FindClass.findAllClass(set, f, "");
		}

		if (set.size() > 0)
			new PredsProfile().profile(set.toArray(new String[set.size()]),
					DirTools.prepare(output_dir + File.separator
							+ DirAndFileConstant.PRED));
	}

	public void profile(String[] args, File outDir) {
		Options.v().set_keep_line_number(true);
		Options.v().setPhaseOption("jb", "use-original-names:true");
		Scene.v().loadBasicClasses();

		HashMap<String, HashMap<String, List<List<Cond>>>> allPreslst = new HashMap<String, HashMap<String, List<List<Cond>>>>();
		for (String strSc : args) {
			SootClass sc = Scene.v().loadClassAndSupport(strSc);

			List<SootMethod> ms = sc.getMethods();
			HashMap<String, List<List<Cond>>> mhm = new HashMap<String, List<List<Cond>>>();

			for (SootMethod sm : ms) {
				if (sm.isAbstract())
					continue;
				Body b = sm.retrieveActiveBody();
				Chain<Unit> units = b.getUnits();
				ExceptionalBlockGraph bGraph = new ExceptionalBlockGraph(b);
				List<Block> blockLs = bGraph.getBlocks();
				List<List<Cond>> predlst = new ArrayList<List<Cond>>();
				Unit u;
				Cond c;
				for (int i = 0; i < blockLs.size(); i++) {
					if ((u = blockLs.get(i).getTail()) instanceof IfStmt) {
						List<Cond> condlst = new ArrayList<Cond>();
						int j = i;
						// c = new Cond(j, Integer.parseInt(u.getTag(
						// "LineNumberTag").toString()), ((IfStmt) u)
						// .getCondition(), findLineOfPred(blockLs, j));

						Unit uu = u;
						LineNumberTag z;
						while ((z = (LineNumberTag) uu.getTag("LineNumberTag")) == null) {
							uu = units.getSuccOf(uu);
						}
						int lnt = z.getLineNumber();
						c = new Cond(j, lnt, ((IfStmt) u).getCondition(), null);
						condlst.add(c);

						while (j < blockLs.size()
								&& (u = blockLs.get(j + 1).getHead()) instanceof IfStmt) {
							List<Block> preBls = blockLs.get(j+1).getPreds();
							if(preBls.size()!=1||preBls.get(0).getIndexInMethod()!=j)
								break;
							j++;

							// c = new Cond(j, Integer.parseInt(u.getTag(
							// "LineNumberTag").toString()), ((IfStmt) u)
							// .getCondition(), findLineOfPred(blockLs, j));

							uu = u;
							while ((z = (LineNumberTag) uu
									.getTag("LineNumberTag")) == null) {
								uu = units.getSuccOf(uu);
							}
							lnt = z.getLineNumber();
							c = new Cond(j, lnt, ((IfStmt) u).getCondition(),
									null);
							condlst.add(c);
						}
						i = j;
						predlst.add(condlst);
					}
				}
				if (predlst.size() > 0) {
					String methodName = sm.toString();
					methodName = methodName.substring(
							methodName.indexOf(":") + 2,
							methodName.length() - 1);
					mhm.put(methodName, predlst);
				}
			}
			if (mhm.keySet().size() > 0)
				allPreslst.put(strSc, mhm);
		}

		// System.out.println();
		for (String className : allPreslst.keySet()) {
			Element xmlRoot = new Element("class");
			xmlRoot.setAttribute("name", className);
			Document xmlDoc = new Document(xmlRoot);
			HashMap<String, List<List<Cond>>> mhm = allPreslst.get(className);

			for (String methodName : mhm.keySet()) {
				List<List<Cond>> predlst = mhm.get(methodName);
				Element xmlMethod = new Element("method");
				xmlMethod.setAttribute("name", methodName);
				for (int i = 0; i < predlst.size(); i++) {
					Element xmlPred = new Element("predicate");
					List<Cond> condlst = predlst.get(i);
					for (int k = 0; k < condlst.size(); k++) {
						Element xmlCond = new Element("condition");
						Cond c = condlst.get(k);
						xmlCond.setAttribute("block", String.valueOf(c.blk));
						xmlCond.setAttribute("line", String.valueOf(c.lineNum));
						xmlCond.setAttribute("cond", c.cond.toString());

						if (c.lineLst != null) {
							Element xmlRelatives = new Element("relatives");
							for (Integer kk : c.lineLst) {
								Element xmlRelate = new Element("relate");
								xmlRelate.setAttribute("line", String
										.valueOf(kk));
								xmlRelatives.addContent(xmlRelate);
							}
							xmlCond.addContent(xmlRelatives);
						}
						xmlPred.addContent(xmlCond);
					}
					xmlMethod.addContent(xmlPred);
				}
				xmlRoot.addContent(xmlMethod);
			}

			try {
				new XMLOutputter().output(xmlDoc, new FileOutputStream(
						new File(outDir, className + ".xml")));
			} catch (FileNotFoundException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}

		}

	}

//	private List<Integer> findLineOfPred(List<Block> blockLs, int i) {
//		Set<Integer> set = new HashSet<Integer>();
//		Iterator<Unit> uIt = blockLs.get(i).iterator();
//		while (uIt.hasNext()) {
//			LineNumberTag z = ((LineNumberTag) uIt.next().getTag(
//					"LineNumberTag"));
//			if (z != null)
//				set.add(z.getLineNumber());
//		}
//		int nxtb = -1;
//
//		for (Block b : blockLs.get(i).getSuccs()) {
//			if (b.getIndexInMethod() != i + 1)
//				nxtb = b.getIndexInMethod();
//		}
//
//		uIt = blockLs.get(nxtb).iterator();
//		while (uIt.hasNext()) {
//			LineNumberTag z = ((LineNumberTag) uIt.next().getTag(
//					"LineNumberTag"));
//			if (z != null)
//				set.add(z.getLineNumber());
//			// it is still right when z==null because under that condition, the
//			// statement's other parts will
//			// add to set.
//		}
//		return new ArrayList<Integer>(set);
//	}

}
